{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Embeddings for Document Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from argparse import Namespace\n",
    "from collections import Counter\n",
    "import json\n",
    "import re\n",
    "import string\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from tqdm import tqdm_notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Vectorization classes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocabulary(object):\n",
    "    \"\"\"Class to process text and extract vocabulary for mapping\"\"\"\n",
    "\n",
    "    def __init__(self, token_to_idx=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            token_to_idx (dict): a pre-existing map of tokens to indices\n",
    "        \"\"\"\n",
    "\n",
    "        if token_to_idx is None:\n",
    "            token_to_idx = {}\n",
    "        self._token_to_idx = token_to_idx\n",
    "\n",
    "        self._idx_to_token = {idx: token \n",
    "                              for token, idx in self._token_to_idx.items()}\n",
    "        \n",
    "    def to_serializable(self):\n",
    "        \"\"\" returns a dictionary that can be serialized \"\"\"\n",
    "        return {'token_to_idx': self._token_to_idx}\n",
    "\n",
    "    @classmethod\n",
    "    def from_serializable(cls, contents):\n",
    "        \"\"\" instantiates the Vocabulary from a serialized dictionary \"\"\"\n",
    "        return cls(**contents)\n",
    "\n",
    "    def add_token(self, token):\n",
    "        \"\"\"Update mapping dicts based on the token.\n",
    "\n",
    "        Args:\n",
    "            token (str): the item to add into the Vocabulary\n",
    "        Returns:\n",
    "            index (int): the integer corresponding to the token\n",
    "        \"\"\"\n",
    "        if token in self._token_to_idx:\n",
    "            index = self._token_to_idx[token]\n",
    "        else:\n",
    "            index = len(self._token_to_idx)\n",
    "            self._token_to_idx[token] = index\n",
    "            self._idx_to_token[index] = token\n",
    "        return index\n",
    "            \n",
    "    def add_many(self, tokens):\n",
    "        \"\"\"Add a list of tokens into the Vocabulary\n",
    "        \n",
    "        Args:\n",
    "            tokens (list): a list of string tokens\n",
    "        Returns:\n",
    "            indices (list): a list of indices corresponding to the tokens\n",
    "        \"\"\"\n",
    "        return [self.add_token(token) for token in tokens]\n",
    "\n",
    "    def lookup_token(self, token):\n",
    "        \"\"\"Retrieve the index associated with the token \n",
    "        \n",
    "        Args:\n",
    "            token (str): the token to look up \n",
    "        Returns:\n",
    "            index (int): the index corresponding to the token\n",
    "        \"\"\"\n",
    "        return self._token_to_idx[token]\n",
    "\n",
    "    def lookup_index(self, index):\n",
    "        \"\"\"Return the token associated with the index\n",
    "        \n",
    "        Args: \n",
    "            index (int): the index to look up\n",
    "        Returns:\n",
    "            token (str): the token corresponding to the index\n",
    "        Raises:\n",
    "            KeyError: if the index is not in the Vocabulary\n",
    "        \"\"\"\n",
    "        if index not in self._idx_to_token:\n",
    "            raise KeyError(\"the index (%d) is not in the Vocabulary\" % index)\n",
    "        return self._idx_to_token[index]\n",
    "\n",
    "    def __str__(self):\n",
    "        return \"<Vocabulary(size=%d)>\" % len(self)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self._token_to_idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SequenceVocabulary(Vocabulary):\n",
    "    def __init__(self, token_to_idx=None, unk_token=\"<UNK>\",\n",
    "                 mask_token=\"<MASK>\", begin_seq_token=\"<BEGIN>\",\n",
    "                 end_seq_token=\"<END>\"):\n",
    "\n",
    "        super(SequenceVocabulary, self).__init__(token_to_idx)\n",
    "\n",
    "        self._mask_token = mask_token\n",
    "        self._unk_token = unk_token\n",
    "        self._begin_seq_token = begin_seq_token\n",
    "        self._end_seq_token = end_seq_token\n",
    "\n",
    "        self.mask_index = self.add_token(self._mask_token)\n",
    "        self.unk_index = self.add_token(self._unk_token)\n",
    "        self.begin_seq_index = self.add_token(self._begin_seq_token)\n",
    "        self.end_seq_index = self.add_token(self._end_seq_token)\n",
    "\n",
    "    def to_serializable(self):\n",
    "        contents = super(SequenceVocabulary, self).to_serializable()\n",
    "        contents.update({'unk_token': self._unk_token,\n",
    "                         'mask_token': self._mask_token,\n",
    "                         'begin_seq_token': self._begin_seq_token,\n",
    "                         'end_seq_token': self._end_seq_token})\n",
    "        return contents\n",
    "\n",
    "    def lookup_token(self, token):\n",
    "        \"\"\"Retrieve the index associated with the token \n",
    "          or the UNK index if token isn't present.\n",
    "        \n",
    "        Args:\n",
    "            token (str): the token to look up \n",
    "        Returns:\n",
    "            index (int): the index corresponding to the token\n",
    "        Notes:\n",
    "            `unk_index` needs to be >=0 (having been added into the Vocabulary) \n",
    "              for the UNK functionality \n",
    "        \"\"\"\n",
    "        if self.unk_index >= 0:\n",
    "            return self._token_to_idx.get(token, self.unk_index)\n",
    "        else:\n",
    "            return self._token_to_idx[token]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Vectorizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewsVectorizer(object):\n",
    "    \"\"\" The Vectorizer which coordinates the Vocabularies and puts them to use\"\"\"    \n",
    "    def __init__(self, title_vocab, category_vocab):\n",
    "        self.title_vocab = title_vocab\n",
    "        self.category_vocab = category_vocab\n",
    "\n",
    "    def vectorize(self, title, vector_length=-1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            title (str): the string of words separated by a space\n",
    "            vector_length (int): an argument for forcing the length of index vector\n",
    "        Returns:\n",
    "            the vetorized title (numpy.array)\n",
    "        \"\"\"\n",
    "        indices = [self.title_vocab.begin_seq_index]\n",
    "        indices.extend(self.title_vocab.lookup_token(token) \n",
    "                       for token in title.split(\" \"))\n",
    "        indices.append(self.title_vocab.end_seq_index)\n",
    "\n",
    "        if vector_length < 0:\n",
    "            vector_length = len(indices)\n",
    "\n",
    "        out_vector = np.zeros(vector_length, dtype=np.int64)\n",
    "        out_vector[:len(indices)] = indices\n",
    "        out_vector[len(indices):] = self.title_vocab.mask_index\n",
    "\n",
    "        return out_vector\n",
    "\n",
    "    @classmethod\n",
    "    def from_dataframe(cls, news_df, cutoff=25):\n",
    "        \"\"\"Instantiate the vectorizer from the dataset dataframe\n",
    "        \n",
    "        Args:\n",
    "            news_df (pandas.DataFrame): the target dataset\n",
    "            cutoff (int): frequency threshold for including in Vocabulary \n",
    "        Returns:\n",
    "            an instance of the NewsVectorizer\n",
    "        \"\"\"\n",
    "        category_vocab = Vocabulary()        \n",
    "        for category in sorted(set(news_df.category)):\n",
    "            category_vocab.add_token(category)\n",
    "\n",
    "        word_counts = Counter()\n",
    "        for title in news_df.title:\n",
    "            for token in title.split(\" \"):\n",
    "                if token not in string.punctuation:\n",
    "                    word_counts[token] += 1\n",
    "        \n",
    "        title_vocab = SequenceVocabulary()\n",
    "        for word, word_count in word_counts.items():\n",
    "            if word_count >= cutoff:\n",
    "                title_vocab.add_token(word)\n",
    "        \n",
    "        return cls(title_vocab, category_vocab)\n",
    "\n",
    "    @classmethod\n",
    "    def from_serializable(cls, contents):\n",
    "        title_vocab = \\\n",
    "            SequenceVocabulary.from_serializable(contents['title_vocab'])\n",
    "        category_vocab =  \\\n",
    "            Vocabulary.from_serializable(contents['category_vocab'])\n",
    "\n",
    "        return cls(title_vocab=title_vocab, category_vocab=category_vocab)\n",
    "\n",
    "    def to_serializable(self):\n",
    "        return {'title_vocab': self.title_vocab.to_serializable(),\n",
    "                'category_vocab': self.category_vocab.to_serializable()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewsDataset(Dataset):\n",
    "    def __init__(self, news_df, vectorizer):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            news_df (pandas.DataFrame): the dataset\n",
    "            vectorizer (NewsVectorizer): vectorizer instatiated from dataset\n",
    "        \"\"\"\n",
    "        self.news_df = news_df\n",
    "        self._vectorizer = vectorizer\n",
    "\n",
    "        # +1 if only using begin_seq, +2 if using both begin and end seq tokens\n",
    "        measure_len = lambda context: len(context.split(\" \"))\n",
    "        self._max_seq_length = max(map(measure_len, news_df.title)) + 2\n",
    "        \n",
    "\n",
    "        self.train_df = self.news_df[self.news_df.split=='train']\n",
    "        self.train_size = len(self.train_df)\n",
    "\n",
    "        self.val_df = self.news_df[self.news_df.split=='val']\n",
    "        self.validation_size = len(self.val_df)\n",
    "\n",
    "        self.test_df = self.news_df[self.news_df.split=='test']\n",
    "        self.test_size = len(self.test_df)\n",
    "\n",
    "        self._lookup_dict = {'train': (self.train_df, self.train_size),\n",
    "                             'val': (self.val_df, self.validation_size),\n",
    "                             'test': (self.test_df, self.test_size)}\n",
    "\n",
    "        self.set_split('train')\n",
    "\n",
    "        # Class weights\n",
    "        class_counts = news_df.category.value_counts().to_dict()\n",
    "        def sort_key(item):\n",
    "            return self._vectorizer.category_vocab.lookup_token(item[0])\n",
    "        sorted_counts = sorted(class_counts.items(), key=sort_key)\n",
    "        frequencies = [count for _, count in sorted_counts]\n",
    "        self.class_weights = 1.0 / torch.tensor(frequencies, dtype=torch.float32)\n",
    "        \n",
    "        \n",
    "    @classmethod\n",
    "    def load_dataset_and_make_vectorizer(cls, news_csv):\n",
    "        \"\"\"Load dataset and make a new vectorizer from scratch\n",
    "        \n",
    "        Args:\n",
    "            surname_csv (str): location of the dataset\n",
    "        Returns:\n",
    "            an instance of SurnameDataset\n",
    "        \"\"\"\n",
    "        news_df = pd.read_csv(news_csv)\n",
    "        train_news_df = news_df[news_df.split=='train']\n",
    "        return cls(news_df, NewsVectorizer.from_dataframe(train_news_df))\n",
    "\n",
    "    @classmethod\n",
    "    def load_dataset_and_load_vectorizer(cls, news_csv, vectorizer_filepath):\n",
    "        \"\"\"Load dataset and the corresponding vectorizer. \n",
    "        Used in the case in the vectorizer has been cached for re-use\n",
    "        \n",
    "        Args:\n",
    "            surname_csv (str): location of the dataset\n",
    "            vectorizer_filepath (str): location of the saved vectorizer\n",
    "        Returns:\n",
    "            an instance of SurnameDataset\n",
    "        \"\"\"\n",
    "        news_df = pd.read_csv(news_csv)\n",
    "        vectorizer = cls.load_vectorizer_only(vectorizer_filepath)\n",
    "        return cls(news_csv, vectorizer)\n",
    "\n",
    "    @staticmethod\n",
    "    def load_vectorizer_only(vectorizer_filepath):\n",
    "        \"\"\"a static method for loading the vectorizer from file\n",
    "        \n",
    "        Args:\n",
    "            vectorizer_filepath (str): the location of the serialized vectorizer\n",
    "        Returns:\n",
    "            an instance of SurnameVectorizer\n",
    "        \"\"\"\n",
    "        with open(vectorizer_filepath) as fp:\n",
    "            return NameVectorizer.from_serializable(json.load(fp))\n",
    "\n",
    "    def save_vectorizer(self, vectorizer_filepath):\n",
    "        \"\"\"saves the vectorizer to disk using json\n",
    "        \n",
    "        Args:\n",
    "            vectorizer_filepath (str): the location to save the vectorizer\n",
    "        \"\"\"\n",
    "        with open(vectorizer_filepath, \"w\") as fp:\n",
    "            json.dump(self._vectorizer.to_serializable(), fp)\n",
    "\n",
    "    def get_vectorizer(self):\n",
    "        \"\"\" returns the vectorizer \"\"\"\n",
    "        return self._vectorizer\n",
    "\n",
    "    def set_split(self, split=\"train\"):\n",
    "        \"\"\" selects the splits in the dataset using a column in the dataframe \"\"\"\n",
    "        self._target_split = split\n",
    "        self._target_df, self._target_size = self._lookup_dict[split]\n",
    "\n",
    "    def __len__(self):\n",
    "        return self._target_size\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"the primary entry point method for PyTorch datasets\n",
    "        \n",
    "        Args:\n",
    "            index (int): the index to the data point \n",
    "        Returns:\n",
    "            a dictionary holding the data point's features (x_data) and label (y_target)\n",
    "        \"\"\"\n",
    "        row = self._target_df.iloc[index]\n",
    "\n",
    "        title_vector = \\\n",
    "            self._vectorizer.vectorize(row.title, self._max_seq_length)\n",
    "\n",
    "        category_index = \\\n",
    "            self._vectorizer.category_vocab.lookup_token(row.category)\n",
    "\n",
    "        return {'x_data': title_vector,\n",
    "                'y_target': category_index}\n",
    "\n",
    "    def get_num_batches(self, batch_size):\n",
    "        \"\"\"Given a batch size, return the number of batches in the dataset\n",
    "        \n",
    "        Args:\n",
    "            batch_size (int)\n",
    "        Returns:\n",
    "            number of batches in the dataset\n",
    "        \"\"\"\n",
    "        return len(self) // batch_size\n",
    "\n",
    "def generate_batches(dataset, batch_size, shuffle=True,\n",
    "                     drop_last=True, device=\"cpu\"): \n",
    "    \"\"\"\n",
    "    A generator function which wraps the PyTorch DataLoader. It will \n",
    "      ensure each tensor is on the write device location.\n",
    "    \"\"\"\n",
    "    dataloader = DataLoader(dataset=dataset, batch_size=batch_size,\n",
    "                            shuffle=shuffle, drop_last=drop_last)\n",
    "\n",
    "    for data_dict in dataloader:\n",
    "        out_data_dict = {}\n",
    "        for name, tensor in data_dict.items():\n",
    "            out_data_dict[name] = data_dict[name].to(device)\n",
    "        yield out_data_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Model: NewsClassifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewsClassifier(nn.Module):\n",
    "    def __init__(self, embedding_size, num_embeddings, num_channels, \n",
    "                 hidden_dim, num_classes, dropout_p, \n",
    "                 pretrained_embeddings=None, padding_idx=0):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            embedding_size (int): size of the embedding vectors\n",
    "            num_embeddings (int): number of embedding vectors\n",
    "            filter_width (int): width of the convolutional kernels\n",
    "            num_channels (int): number of convolutional kernels per layer\n",
    "            hidden_dim (int): the size of the hidden dimension\n",
    "            num_classes (int): the number of classes in classification\n",
    "            dropout_p (float): a dropout parameter \n",
    "            pretrained_embeddings (numpy.array): previously trained word embeddings\n",
    "                default is None. If provided, \n",
    "            padding_idx (int): an index representing a null position\n",
    "        \"\"\"\n",
    "        super(NewsClassifier, self).__init__()\n",
    "\n",
    "        if pretrained_embeddings is None:\n",
    "\n",
    "            self.emb = nn.Embedding(embedding_dim=embedding_size,\n",
    "                                    num_embeddings=num_embeddings,\n",
    "                                    padding_idx=padding_idx)        \n",
    "        else:\n",
    "            pretrained_embeddings = torch.from_numpy(pretrained_embeddings).float()\n",
    "            self.emb = nn.Embedding(embedding_dim=embedding_size,\n",
    "                                    num_embeddings=num_embeddings,\n",
    "                                    padding_idx=padding_idx,\n",
    "                                    _weight=pretrained_embeddings)\n",
    "        \n",
    "            \n",
    "        self.convnet = nn.Sequential(\n",
    "            nn.Conv1d(in_channels=embedding_size, \n",
    "                   out_channels=num_channels, kernel_size=3),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(in_channels=num_channels, out_channels=num_channels, \n",
    "                   kernel_size=3, stride=2),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(in_channels=num_channels, out_channels=num_channels, \n",
    "                   kernel_size=3, stride=2),\n",
    "            nn.ELU(),\n",
    "            nn.Conv1d(in_channels=num_channels, out_channels=num_channels, \n",
    "                   kernel_size=3),\n",
    "            nn.ELU()\n",
    "        )\n",
    "\n",
    "        self._dropout_p = dropout_p\n",
    "        self.fc1 = nn.Linear(num_channels, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, num_classes)\n",
    "\n",
    "    def forward(self, x_in, apply_softmax=False):\n",
    "        \"\"\"The forward pass of the classifier\n",
    "        \n",
    "        Args:\n",
    "            x_in (torch.Tensor): an input data tensor. \n",
    "                x_in.shape should be (batch, dataset._max_seq_length)\n",
    "            apply_softmax (bool): a flag for the softmax activation\n",
    "                should be false if used with the Cross Entropy losses\n",
    "        Returns:\n",
    "            the resulting tensor. tensor.shape should be (batch, num_classes)\n",
    "        \"\"\"\n",
    "        \n",
    "        # embed and permute so features are channels\n",
    "        x_embedded = self.emb(x_in).permute(0, 2, 1)\n",
    "\n",
    "        features = self.convnet(x_embedded)\n",
    "\n",
    "        # average and remove the extra dimension\n",
    "        remaining_size = features.size(dim=2)\n",
    "        features = F.avg_pool1d(features, remaining_size).squeeze(dim=2)\n",
    "        features = F.dropout(features, p=self._dropout_p)\n",
    "        \n",
    "        # mlp classifier\n",
    "        intermediate_vector = F.relu(F.dropout(self.fc1(features), p=self._dropout_p))\n",
    "        prediction_vector = self.fc2(intermediate_vector)\n",
    "\n",
    "        if apply_softmax:\n",
    "            prediction_vector = F.softmax(prediction_vector, dim=1)\n",
    "\n",
    "        return prediction_vector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Routine"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_train_state(args):\n",
    "    return {'stop_early': False,\n",
    "            'early_stopping_step': 0,\n",
    "            'early_stopping_best_val': 1e8,\n",
    "            'learning_rate': args.learning_rate,\n",
    "            'epoch_index': 0,\n",
    "            'train_loss': [],\n",
    "            'train_acc': [],\n",
    "            'val_loss': [],\n",
    "            'val_acc': [],\n",
    "            'test_loss': -1,\n",
    "            'test_acc': -1,\n",
    "            'model_filename': args.model_state_file}\n",
    "\n",
    "def update_train_state(args, model, train_state):\n",
    "    \"\"\"Handle the training state updates.\n",
    "\n",
    "    Components:\n",
    "     - Early Stopping: Prevent overfitting.\n",
    "     - Model Checkpoint: Model is saved if the model is better\n",
    "\n",
    "    :param args: main arguments\n",
    "    :param model: model to train\n",
    "    :param train_state: a dictionary representing the training state values\n",
    "    :returns:\n",
    "        a new train_state\n",
    "    \"\"\"\n",
    "\n",
    "    # Save one model at least\n",
    "    if train_state['epoch_index'] == 0:\n",
    "        torch.save(model.state_dict(), train_state['model_filename'])\n",
    "        train_state['stop_early'] = False\n",
    "\n",
    "    # Save model if performance improved\n",
    "    elif train_state['epoch_index'] >= 1:\n",
    "        loss_tm1, loss_t = train_state['val_loss'][-2:]\n",
    "\n",
    "        # If loss worsened\n",
    "        if loss_t >= train_state['early_stopping_best_val']:\n",
    "            # Update step\n",
    "            train_state['early_stopping_step'] += 1\n",
    "        # Loss decreased\n",
    "        else:\n",
    "            # Save the best model\n",
    "            if loss_t < train_state['early_stopping_best_val']:\n",
    "                torch.save(model.state_dict(), train_state['model_filename'])\n",
    "\n",
    "            # Reset early stopping step\n",
    "            train_state['early_stopping_step'] = 0\n",
    "\n",
    "        # Stop early ?\n",
    "        train_state['stop_early'] = \\\n",
    "            train_state['early_stopping_step'] >= args.early_stopping_criteria\n",
    "\n",
    "    return train_state\n",
    "\n",
    "def compute_accuracy(y_pred, y_target):\n",
    "    _, y_pred_indices = y_pred.max(dim=1)\n",
    "    n_correct = torch.eq(y_pred_indices, y_target).sum().item()\n",
    "    return n_correct / len(y_pred_indices) * 100"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### general utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_seed_everywhere(seed, cuda):\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    if cuda:\n",
    "        torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "def handle_dirs(dirpath):\n",
    "    if not os.path.exists(dirpath):\n",
    "        os.makedirs(dirpath)\n",
    "        \n",
    "def load_glove_from_file(glove_filepath):\n",
    "    \"\"\"\n",
    "    Load the GloVe embeddings \n",
    "    \n",
    "    Args:\n",
    "        glove_filepath (str): path to the glove embeddings file \n",
    "    Returns:\n",
    "        word_to_index (dict), embeddings (numpy.ndarary)\n",
    "    \"\"\"\n",
    "\n",
    "    word_to_index = {}\n",
    "    embeddings = []\n",
    "    with open(glove_filepath, \"r\") as fp:\n",
    "        for index, line in enumerate(fp):\n",
    "            line = line.split(\" \") # each line: word num1 num2 ...\n",
    "            word_to_index[line[0]] = index # word = line[0] \n",
    "            embedding_i = np.array([float(val) for val in line[1:]])\n",
    "            embeddings.append(embedding_i)\n",
    "    return word_to_index, np.stack(embeddings)\n",
    "\n",
    "def make_embedding_matrix(glove_filepath, words):\n",
    "    \"\"\"\n",
    "    Create embedding matrix for a specific set of words.\n",
    "    \n",
    "    Args:\n",
    "        glove_filepath (str): file path to the glove embeddigns\n",
    "        words (list): list of words in the dataset\n",
    "    \"\"\"\n",
    "    word_to_idx, glove_embeddings = load_glove_from_file(glove_filepath)\n",
    "    embedding_size = glove_embeddings.shape[1]\n",
    "    \n",
    "    final_embeddings = np.zeros((len(words), embedding_size))\n",
    "\n",
    "    for i, word in enumerate(words):\n",
    "        if word in word_to_idx:\n",
    "            final_embeddings[i, :] = glove_embeddings[word_to_idx[word]]\n",
    "        else:\n",
    "            embedding_i = torch.ones(1, embedding_size)\n",
    "            torch.nn.init.xavier_uniform_(embedding_i)\n",
    "            final_embeddings[i, :] = embedding_i\n",
    "\n",
    "    return final_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings and some prep work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from argparse import Namespace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expanded filepaths: \n",
      "\tmodel_storage/ch5/document_classification/vectorizer.json\n",
      "\tmodel_storage/ch5/document_classification/model.pth\n",
      "Using CUDA: True\n"
     ]
    }
   ],
   "source": [
    "args = Namespace(\n",
    "    # Data and Path hyper parameters\n",
    "    news_csv=\"data/ag_news/news_with_splits.csv\",\n",
    "    vectorizer_file=\"vectorizer.json\",\n",
    "    model_state_file=\"model.pth\",\n",
    "    save_dir=\"model_storage/ch5/document_classification\",\n",
    "    # Model hyper parameters\n",
    "    glove_filepath='data/glove/glove.6B.100d.txt', \n",
    "    use_glove=False,\n",
    "    embedding_size=100, \n",
    "    hidden_dim=100, \n",
    "    num_channels=100, \n",
    "    # Training hyper parameter\n",
    "    seed=1337, \n",
    "    learning_rate=0.001, \n",
    "    dropout_p=0.1, \n",
    "    batch_size=128, \n",
    "    num_epochs=100, \n",
    "    early_stopping_criteria=5, \n",
    "    # Runtime option\n",
    "    cuda=True, \n",
    "    catch_keyboard_interrupt=True, \n",
    "    reload_from_files=False,\n",
    "    expand_filepaths_to_save_dir=True\n",
    ") \n",
    "\n",
    "if args.expand_filepaths_to_save_dir:\n",
    "    args.vectorizer_file = os.path.join(args.save_dir,\n",
    "                                        args.vectorizer_file)\n",
    "\n",
    "    args.model_state_file = os.path.join(args.save_dir,\n",
    "                                         args.model_state_file)\n",
    "    \n",
    "    print(\"Expanded filepaths: \")\n",
    "    print(\"\\t{}\".format(args.vectorizer_file))\n",
    "    print(\"\\t{}\".format(args.model_state_file))\n",
    "    \n",
    "# Check CUDA\n",
    "if not torch.cuda.is_available():\n",
    "    args.cuda = False\n",
    "    \n",
    "args.device = torch.device(\"cuda\" if args.cuda else \"cpu\")\n",
    "print(\"Using CUDA: {}\".format(args.cuda))\n",
    "\n",
    "# Set seed for reproducibility\n",
    "set_seed_everywhere(args.seed, args.cuda)\n",
    "\n",
    "# handle dirs\n",
    "handle_dirs(args.save_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "args.use_glove = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using pre-trained embeddings\n"
     ]
    }
   ],
   "source": [
    "if args.reload_from_files:\n",
    "    # training from a checkpoint\n",
    "    dataset = NewsDataset.load_dataset_and_load_vectorizer(args.news_csv,\n",
    "                                                              args.vectorizer_file)\n",
    "else:\n",
    "    # create dataset and vectorizer\n",
    "    dataset = NewsDataset.load_dataset_and_make_vectorizer(args.news_csv)\n",
    "    dataset.save_vectorizer(args.vectorizer_file)\n",
    "vectorizer = dataset.get_vectorizer()\n",
    "\n",
    "# Use GloVe or randomly initialized embeddings\n",
    "if args.use_glove:\n",
    "    words = vectorizer.title_vocab._token_to_idx.keys()\n",
    "    embeddings = make_embedding_matrix(glove_filepath=args.glove_filepath, \n",
    "                                       words=words)\n",
    "    print(\"Using pre-trained embeddings\")\n",
    "else:\n",
    "    print(\"Not using pre-trained embeddings\")\n",
    "    embeddings = None\n",
    "\n",
    "classifier = NewsClassifier(embedding_size=args.embedding_size, \n",
    "                            num_embeddings=len(vectorizer.title_vocab),\n",
    "                            num_channels=args.num_channels,\n",
    "                            hidden_dim=args.hidden_dim, \n",
    "                            num_classes=len(vectorizer.category_vocab), \n",
    "                            dropout_p=args.dropout_p,\n",
    "                            pretrained_embeddings=embeddings,\n",
    "                            padding_idx=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7387211e460e433d8cf58de04a721383",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='training routine', style=ProgressStyle(description_width='ini…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a71b16b317474357b4ecf26ea34ba7a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='split=train', max=656, style=ProgressStyle(description_width=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7065e351bda7438cb73588159eb0f6ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, description='split=val', max=140, style=ProgressStyle(description_width='i…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exiting loop\n"
     ]
    }
   ],
   "source": [
    "classifier = classifier.to(args.device)\n",
    "dataset.class_weights = dataset.class_weights.to(args.device)\n",
    "    \n",
    "loss_func = nn.CrossEntropyLoss(dataset.class_weights)\n",
    "optimizer = optim.Adam(classifier.parameters(), lr=args.learning_rate)\n",
    "scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=optimizer,\n",
    "                                           mode='min', factor=0.5,\n",
    "                                           patience=1)\n",
    "\n",
    "train_state = make_train_state(args)\n",
    "\n",
    "epoch_bar = tqdm_notebook(desc='training routine', \n",
    "                          total=args.num_epochs,\n",
    "                          position=0)\n",
    "\n",
    "dataset.set_split('train')\n",
    "train_bar = tqdm_notebook(desc='split=train',\n",
    "                          total=dataset.get_num_batches(args.batch_size), \n",
    "                          position=1, \n",
    "                          leave=True)\n",
    "dataset.set_split('val')\n",
    "val_bar = tqdm_notebook(desc='split=val',\n",
    "                        total=dataset.get_num_batches(args.batch_size), \n",
    "                        position=1, \n",
    "                        leave=True)\n",
    "\n",
    "try:\n",
    "    for epoch_index in range(args.num_epochs):\n",
    "        train_state['epoch_index'] = epoch_index\n",
    "\n",
    "        # Iterate over training dataset\n",
    "\n",
    "        # setup: batch generator, set loss and acc to 0, set train mode on\n",
    "\n",
    "        dataset.set_split('train')\n",
    "        batch_generator = generate_batches(dataset, \n",
    "                                           batch_size=args.batch_size, \n",
    "                                           device=args.device)\n",
    "        running_loss = 0.0\n",
    "        running_acc = 0.0\n",
    "        classifier.train()\n",
    "\n",
    "        for batch_index, batch_dict in enumerate(batch_generator):\n",
    "            # the training routine is these 5 steps:\n",
    "\n",
    "            # --------------------------------------\n",
    "            # step 1. zero the gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # step 2. compute the output\n",
    "            y_pred = classifier(batch_dict['x_data'])\n",
    "\n",
    "            # step 3. compute the loss\n",
    "            loss = loss_func(y_pred, batch_dict['y_target'])\n",
    "            loss_t = loss.item()\n",
    "            running_loss += (loss_t - running_loss) / (batch_index + 1)\n",
    "\n",
    "            # step 4. use loss to produce gradients\n",
    "            loss.backward()\n",
    "\n",
    "            # step 5. use optimizer to take gradient step\n",
    "            optimizer.step()\n",
    "            # -----------------------------------------\n",
    "            # compute the accuracy\n",
    "            acc_t = compute_accuracy(y_pred, batch_dict['y_target'])\n",
    "            running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "\n",
    "            # update bar\n",
    "            train_bar.set_postfix(loss=running_loss, acc=running_acc, \n",
    "                                  epoch=epoch_index)\n",
    "            train_bar.update()\n",
    "\n",
    "        train_state['train_loss'].append(running_loss)\n",
    "        train_state['train_acc'].append(running_acc)\n",
    "\n",
    "        # Iterate over val dataset\n",
    "\n",
    "        # setup: batch generator, set loss and acc to 0; set eval mode on\n",
    "        dataset.set_split('val')\n",
    "        batch_generator = generate_batches(dataset, \n",
    "                                           batch_size=args.batch_size, \n",
    "                                           device=args.device)\n",
    "        running_loss = 0.\n",
    "        running_acc = 0.\n",
    "        classifier.eval()\n",
    "\n",
    "        for batch_index, batch_dict in enumerate(batch_generator):\n",
    "\n",
    "            # compute the output\n",
    "            y_pred =  classifier(batch_dict['x_data'])\n",
    "\n",
    "            # step 3. compute the loss\n",
    "            loss = loss_func(y_pred, batch_dict['y_target'])\n",
    "            loss_t = loss.item()\n",
    "            running_loss += (loss_t - running_loss) / (batch_index + 1)\n",
    "\n",
    "            # compute the accuracy\n",
    "            acc_t = compute_accuracy(y_pred, batch_dict['y_target'])\n",
    "            running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "            val_bar.set_postfix(loss=running_loss, acc=running_acc, \n",
    "                            epoch=epoch_index)\n",
    "            val_bar.update()\n",
    "\n",
    "        train_state['val_loss'].append(running_loss)\n",
    "        train_state['val_acc'].append(running_acc)\n",
    "\n",
    "        train_state = update_train_state(args=args, model=classifier,\n",
    "                                         train_state=train_state)\n",
    "\n",
    "        scheduler.step(train_state['val_loss'][-1])\n",
    "\n",
    "        if train_state['stop_early']:\n",
    "            break\n",
    "\n",
    "        train_bar.n = 0\n",
    "        val_bar.n = 0\n",
    "        epoch_bar.update()\n",
    "except KeyboardInterrupt:\n",
    "    print(\"Exiting loop\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the loss & accuracy on the test set using the best available model\n",
    "\n",
    "classifier.load_state_dict(torch.load(train_state['model_filename']))\n",
    "\n",
    "classifier = classifier.to(args.device)\n",
    "dataset.class_weights = dataset.class_weights.to(args.device)\n",
    "loss_func = nn.CrossEntropyLoss(dataset.class_weights)\n",
    "\n",
    "dataset.set_split('test')\n",
    "batch_generator = generate_batches(dataset, \n",
    "                                   batch_size=args.batch_size, \n",
    "                                   device=args.device)\n",
    "running_loss = 0.\n",
    "running_acc = 0.\n",
    "classifier.eval()\n",
    "\n",
    "for batch_index, batch_dict in enumerate(batch_generator):\n",
    "    # compute the output\n",
    "    y_pred =  classifier(batch_dict['x_data'])\n",
    "    \n",
    "    # compute the loss\n",
    "    loss = loss_func(y_pred, batch_dict['y_target'])\n",
    "    loss_t = loss.item()\n",
    "    running_loss += (loss_t - running_loss) / (batch_index + 1)\n",
    "\n",
    "    # compute the accuracy\n",
    "    acc_t = compute_accuracy(y_pred, batch_dict['y_target'])\n",
    "    running_acc += (acc_t - running_acc) / (batch_index + 1)\n",
    "\n",
    "train_state['test_loss'] = running_loss\n",
    "train_state['test_acc'] = running_acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 0.45378321813685574;\n",
      "Test Accuracy: 83.77790178571429\n"
     ]
    }
   ],
   "source": [
    "print(\"Test loss: {};\".format(train_state['test_loss']))\n",
    "print(\"Test Accuracy: {}\".format(train_state['test_acc']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preprocess the reviews\n",
    "def preprocess_text(text):\n",
    "    text = ' '.join(word.lower() for word in text.split(\" \"))\n",
    "    text = re.sub(r\"([.,!?])\", r\" \\1 \", text)\n",
    "    text = re.sub(r\"[^a-zA-Z.,!?]+\", r\" \", text)\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_category(title, classifier, vectorizer, max_length):\n",
    "    \"\"\"Predict a News category for a new title\n",
    "    \n",
    "    Args:\n",
    "        title (str): a raw title string\n",
    "        classifier (NewsClassifier): an instance of the trained classifier\n",
    "        vectorizer (NewsVectorizer): the corresponding vectorizer\n",
    "        max_length (int): the max sequence length\n",
    "            Note: CNNs are sensitive to the input data tensor size. \n",
    "                  This ensures to keep it the same size as the training data\n",
    "    \"\"\"\n",
    "    title = preprocess_text(title)\n",
    "    vectorized_title = \\\n",
    "        torch.tensor(vectorizer.vectorize(title, vector_length=max_length))\n",
    "    result = classifier(vectorized_title.unsqueeze(0), apply_softmax=True)\n",
    "    probability_values, indices = result.max(dim=1)\n",
    "    predicted_category = vectorizer.category_vocab.lookup_index(indices.item())\n",
    "\n",
    "    return {'category': predicted_category, \n",
    "            'probability': probability_values.item()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_samples():\n",
    "    samples = {}\n",
    "    for cat in dataset.val_df.category.unique():\n",
    "        samples[cat] = dataset.val_df.title[dataset.val_df.category==cat].tolist()[:5]\n",
    "    return samples\n",
    "\n",
    "val_samples = get_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True Category: Business\n",
      "==============================\n",
      "Prediction: Business (p=0.77)\n",
      "\t + Sample: az suspends marketing of cancer drug\n",
      "Prediction: Business (p=0.75)\n",
      "\t + Sample: business world has mixed reaction to perez move\n",
      "Prediction: Sports (p=0.68)\n",
      "\t + Sample: betting against bombay\n",
      "Prediction: Business (p=0.94)\n",
      "\t + Sample: malpractice insurers face a tough market\n",
      "Prediction: Sci/Tech (p=0.80)\n",
      "\t + Sample: nvidia is vindicated\n",
      "------------------------------\n",
      "\n",
      "True Category: Sci/Tech\n",
      "==============================\n",
      "Prediction: World (p=0.52)\n",
      "\t + Sample: spies prize webcam s eyes\n",
      "Prediction: Sci/Tech (p=0.97)\n",
      "\t + Sample: sober worm causes headaches\n",
      "Prediction: Sci/Tech (p=0.75)\n",
      "\t + Sample: local search missing pieces falling into place\n",
      "Prediction: Sci/Tech (p=0.92)\n",
      "\t + Sample: hackers baiting internet users with beckham pix\n",
      "Prediction: Sci/Tech (p=0.96)\n",
      "\t + Sample: nokia adds blackberry support to series handsets\n",
      "------------------------------\n",
      "\n",
      "True Category: Sports\n",
      "==============================\n",
      "Prediction: Sports (p=0.85)\n",
      "\t + Sample: is meyer the man to get irish up ? \n",
      "Prediction: Sports (p=0.76)\n",
      "\t + Sample: who ? who ? and clemens\n",
      "Prediction: Sports (p=0.96)\n",
      "\t + Sample: baseball today ap \n",
      "Prediction: Sports (p=0.85)\n",
      "\t + Sample: mark kreidler yao ming epitomizes the chinese athlete who is lt b gt . . . lt b gt \n",
      "Prediction: Sports (p=0.96)\n",
      "\t + Sample: no . miami rebounds to beat fsu in overtime\n",
      "------------------------------\n",
      "\n",
      "True Category: World\n",
      "==============================\n",
      "Prediction: World (p=0.97)\n",
      "\t + Sample: arafat in pain but expected to recover shaath\n",
      "Prediction: World (p=0.97)\n",
      "\t + Sample: maoist rebels bomb kathmandu building , no injuries reuters \n",
      "Prediction: World (p=1.00)\n",
      "\t + Sample: son running for ill . rep . s house seat ap \n",
      "Prediction: World (p=0.91)\n",
      "\t + Sample: strong quake hits in japan\n",
      "Prediction: World (p=1.00)\n",
      "\t + Sample: israel assassinates hamas militant in damascus\n",
      "------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#title = input(\"Enter a news title to classify: \")\n",
    "classifier = classifier.to(\"cpu\")\n",
    "\n",
    "for truth, sample_group in val_samples.items():\n",
    "    print(f\"True Category: {truth}\")\n",
    "    print(\"=\"*30)\n",
    "    for sample in sample_group:\n",
    "        prediction = predict_category(sample, classifier, \n",
    "                                      vectorizer, dataset._max_seq_length + 1)\n",
    "        print(\"Prediction: {} (p={:0.2f})\".format(prediction['category'],\n",
    "                                                  prediction['probability']))\n",
    "        print(\"\\t + Sample: {}\".format(sample))\n",
    "    print(\"-\"*30 + \"\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlpbook",
   "language": "python",
   "name": "nlpbook"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.2"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "138px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "threshold": "5",
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
